﻿#ifndef RBTree_h__
#define RBTree_h__

/********************************************************************
	创建于:		2011/05/06
	创建于:		6:5:2011   11:38
	基本文件:	RBTree
	文件扩展:	h
	作者:		王凯
	
	目的:		使用c++对算法导论中的红黑二叉树的实现。
*********************************************************************/

#include <iostream>

template<typename KEY, typename OBJECT>
class RBTree
{
public:
	enum RBColor{RED, BLACK};
	struct RBNode
	{
		RBNode* parent;
		RBNode* left;
		RBNode* right;
		RBColor color;
		OBJECT data;
		KEY key;

		RBNode()
		{
			parent = NULL;
			left = NULL;
			right = NULL;
			color = RED;
		}
		//对节点的输出符重载
		friend std::ostream& operator<<(std::ostream& out,const RBNode* n);
	};
	//对节点的输出符重载
	friend std::ostream& operator<<(std::ostream& out,const RBNode* n)
	{

		out<<"Data: "<<n->data<<" Color: ";
		if(n->color == BLACK)
			out<<"Black"<<std::endl;
		else
			out<<"Red"<<std::endl;

		return out;
	}

	RBTree()
	{
		m_nullNode = new RBNode();
		
		m_Root = m_nullNode;

		m_nullNode->left = m_Root;
		m_nullNode->right = m_Root;
		m_nullNode->parent = m_Root;
		m_nullNode->color = BLACK;
	}

	
	~RBTree()
	{
		clear(m_Root);
		delete m_nullNode;
	}

	void printTree(std::ostream& out = std::cout)
	{
		printTree(m_Root, out);
	}

	RBNode* minimum(RBNode* t)
	{
		while(t->left != m_nullNode)
			t = t->left;

		return t;
	}

	//插入一个值到树中
	bool insert(const KEY key, const OBJECT z)
	{
		RBNode* y = m_nullNode;
		RBNode* x = m_Root;

		while(x != m_nullNode)
		{
			y = x;
			if(key < x->key)
				x = x->left;
			else if(key > x->key)
				x = x->right;
			else
				return false;
		}

		RBNode* newNode = new RBNode();
		newNode->data = z;
		newNode->key = key;
		newNode->parent = y;
		newNode->left = m_nullNode;
		newNode->right = m_nullNode;
		newNode->color = RED;

		if(y == m_nullNode)
		{
			m_Root = newNode;
			m_Root->parent = m_nullNode;
			m_nullNode->left = m_Root;
			m_nullNode->right = m_Root;
			m_nullNode->parent = m_Root;
		}
		else
		{
			if(key < y->key)
				y->left = newNode;
			else
				y->right = newNode;
		}
		
		insertFixup(newNode);

		return true;
	}

	//从树中删除一个节点（如果有）
	//找到并删除返回真
	//没找到返回假
	bool remove(const KEY key)
	{
		//返回key所对应的节点指针
		RBNode* z = findNode(key);
		//如果没有找到，则返回假
		if(z == m_nullNode)
			return false;

		//y用来保存待删节点（z）的子节点
		RBNode* y;	

		//找到后，进一步确认删除步骤
		//如果需要删除的节点有子节点，则把z赋值给y
		//不存在则把z的后继赋值给y
		if(z->left == m_nullNode || z->right == m_nullNode)
		{
				y = z;
		}	
		else
		{
			RBNode* suc = successor(z);
			z->data = suc->data;
			z->key = suc->key;
			z = suc;
		}
	
		//让y的一个有效子树赋值给x，如果没有任何子树，则被赋值为空节点
		if(z->left != m_nullNode)
			y = z->left;
		else
			y = z->right;

		//更新y的父节点指针
		y->parent = z->parent;
		//如果z是根节点，则把y指定为根节点
		if(z->parent == m_nullNode)
		{
			m_Root = y;
			m_nullNode->parent = m_Root;
			m_nullNode->left = m_Root;
			m_nullNode->right = m_Root;
		}
		//z不是根节点，且z是其父亲的右儿子
		else if(z->parent->right == z)
			z->parent->right = y;
		//z不是根节点，且z只可能是其父亲的左儿子
		else
			z->parent->left = y;

		//如果删除的节点颜色是黑色，则肯定会破坏红黑树性质，需要进行修补
		if(z->color==BLACK)
			removeFixup(y);

		delete z;
		return true;
	}

	RBNode* findNode(const KEY& k)
	{
		//return findNode(m_Root, z);
		RBNode* index = m_Root;
		while(index != m_nullNode)
		{
			if(k < index->key)
				index = index->left;
			else if(k > index->key)
				index = index->right;
			else
				break;
		}

		return index;
	}
private:
	RBNode* m_Root;
	RBNode* m_nullNode;

	//先序遍历输出每个节点
	void printTree(RBNode* t, std::ostream& out = std::cout)
	{
		if(t != m_nullNode)
		{
			

 			out<<"Data： "<<t->data<<"， Key： "<<t->key<<"， Color:：";
 			if(t->color == RED)
 				out<<"Color： RED";
 			else
 				out<<"Color：BLACK";
 			out<<" Parent:："<<t->parent->data<<std::endl;

			if(t->left != m_nullNode)
				printTree(t->left, out );
			if(t->right != m_nullNode)
				printTree(t->right, out);
		}
	}

	//找寻节点
	RBNode* findNode(RBNode* t, const OBJECT& z)
	{
		if(t == m_nullNode)
			return m_nullNode;
		else
		{
			if(z < t->data)
				return findNode(t->left, z);
			else if(z > t->data)
				return findNode(t->right, z);
			else
				return t;
		}
	}

	//左旋节点
	void left_Rotate(RBNode* x)
	{
		if(x == m_nullNode || x->right == m_nullNode)
			return;
		RBNode* y = x->right;
		x->right = y->left;
		if(y->left != m_nullNode)
			y->left->parent = x;
		y->parent = x->parent;
		if(x->parent == m_nullNode)
		{
			m_Root = y;
			m_nullNode->left = m_Root;
			m_nullNode->right = m_Root;
		}
		else 
		{
			if(x == x->parent->left)
				x->parent->left = y;
			else
				x->parent->right = y;
		}
		y->left = x;
		x->parent = y;
	}

	//右旋节点
	void right_Rotate(RBNode* y)
	{
		if(y == m_nullNode || y == m_nullNode)
			return;
		RBNode* x = y->left;
		y->left = x->right;
		if(x->right != m_nullNode)
			x->right->parent = y;
		x->parent = y->parent;
		if(y->parent == m_nullNode)
		{
			m_Root = x;
			m_nullNode->left = m_Root;
			m_nullNode->right = m_Root;
		}
		else
		{
			if(y == y->parent->right)
				y->parent->right = x;
			else
				y->parent->left = x;
		}
		y->parent = x;
		x->right = y;
	}

	//插入值后，可能导致红黑树的2、5性质被破坏，需要用该例程来进行调整。
	void insertFixup(RBNode* z)
	{
		while(z->parent->color == RED)
		{
			if(z->parent == z->parent->parent->left)//如果z的父节点是z的爷爷节点的左儿子
			{
				RBNode* y = z->parent->parent->right;
				if(y->color == RED)//case 1
				{
					z->parent->color = BLACK;
					y->color = BLACK;
					z->parent->parent->color = RED;
					z = z->parent->parent;
				}
				else// y->color == BLACK
				{
					if(z == z->parent->right)//case 2
					{
						z = z->parent;
						left_Rotate(z);
					}
					else//case 3
					{
						z->parent->color = BLACK;
						z->parent->parent->color = RED;
						right_Rotate(z->parent->parent);
					}
				}
			}
			else//如果z的父节点是z的爷爷节点的右儿子
			{
				RBNode* y = z->parent->parent->left;
				if(y->color == RED)
				{
					z->parent->color = BLACK;
					y->color = BLACK;
					z->parent->parent->color = RED;
					z = z->parent->parent;
				}
				else// y->color == BLACK
				{
					if(z == z->parent->left)//case 2
					{
						z = z->parent;
						right_Rotate(z);
					}
					else
					{
						z->parent->color = BLACK;
						z->parent->parent->color = RED;
						left_Rotate(z->parent->parent);
					}
				}
			}
		}
		m_Root->color = BLACK;
	}

	//返回x的后继
	RBNode* successor(RBNode* x)
	{
		//空节点没有后继
		if(x == m_nullNode)
			return m_nullNode;

		RBNode* y = x->right;

		//x的右子树存在，则找到右子树最小节点（即右子树最左节点）
 		while(y != m_nullNode)
 		{
 			if(y->left != m_nullNode)
 				y = y->left;
 			else
 				break;
 		}

		//x的右子树不存在或者右子树的左子树不存在
		if(y == m_nullNode)
		{
			RBNode* p = y->parent;
			y = x;
			//找一个节点（其是父节点的左儿子）
			while(p != m_nullNode && y == p->right)
			{
				y = p;
				p = p->parent;
			}
			y = p;
		}
		return y;
	}

	//对删除节点后出现的不平衡情况经行修复
	void removeFixup(RBNode* x)
	{
		//w是x的兄弟节点
		RBNode* w = m_nullNode;
		while(x != m_Root && x->color == BLACK)
		{
			//x是其父节点的左儿子
			if(x == x->parent->left)
			{
				//将x的兄弟节点赋给w
				w = x->parent->right;

				//w的颜色是红色
				if(w->color == RED)//case 1
				{
					//这时需要将w的颜色改为黑色，并将x的父节点颜色变为红色
					//旋转x的父节点之后，更新w的指针
					w->color = BLACK;
					x->parent->color = RED;
					left_Rotate(x->parent);
					w = x->parent->right;
				}
				//w的颜色是黑色
				else
				{
					//且w的两个子节点为黑色
					if(w->left->color == BLACK && w->right->color == BLACK)//case 2
					{
						//这时只需要将w的颜色改为红色，更新x为x的父节点即可
						w->color = RED;
						x = x->parent;
					}
					//w的右孩子为黑色
					else 
					{
						if(w->right->color == BLACK)//case 3
						{
							//将case3转换为case4加以处理
							//改变w的左孩子颜色为黑色，改变w的颜色为红色
							//右旋w，更新w的指针
							w->left->color = BLACK;	
							w->color = RED;
							right_Rotate(w);
							w = x->parent->right;
						}
						//w的孩子为红色
						//case 4
						w->color = x->parent->color;
						x->parent->color = BLACK;
						w->right->color = BLACK;
						left_Rotate(x->parent);
						x = m_Root;
					}
				}
			}
			//x是其父节点的右儿子
			else
			{
				w = x->parent->left;

				//w的颜色是红色
				if(w->color == RED)//case 1
				{
					//这时需要将w的颜色改为黑色，并将x的父节点颜色变为红色
					//旋转x的父节点之后，更新w的指针
					w->color = BLACK;
					x->parent->color = RED;
					right_Rotate(x->parent);
					w = x->parent->left;
				}
				//w的颜色是黑色
				else
				{
					//且w的两个子节点为黑色
					if(w->right->color == BLACK && w->left->color == BLACK)//case 2
					{
						//这时只需要将w的颜色改为红色，更新x为x的父节点即可
						w->color = RED;
						x = x->parent;
					}
					//w的右孩子为黑色
					else 
					{
						if(w->left->color == BLACK)//case 3
						{
							//将case3转换为case4加以处理
							//改变w的左孩子颜色为黑色，改变w的颜色为红色
							//右旋w，更新w的指针
							w->right->color = BLACK;
							w->color = RED;
							left_Rotate(w);
							w = w->parent->left;
						}
						//w的孩子为红色
						w->color = x->parent->color;
						x->parent->color = BLACK;
						w->left->color = BLACK;
						right_Rotate(x->parent);
						x = m_Root;
					}
				}
			}
		}
		x->color = BLACK;
	}

	void clear(RBNode* p)
	{
		if(p == m_nullNode)
			return ;
		clear(p->left);
		clear(p->right);
		delete p;
	}

};

#endif // RBTree_h__
